from transformers import BertTokenizer

from dataProcess import train_text, train_label, val_text, val_label, test_text, test_label

tokenizer = BertTokenizer.from_pretrained('./bert-base-chinese')

train_token = tokenizer(train_text, truncation=True, padding='max_length', max_length=512)
val_token = tokenizer(val_text, truncation=True, padding='max_length', max_length=512)
test_token = tokenizer(test_text, truncation=True, padding='max_length', max_length=512)

# print(train_token['input_ids'])

import torch


class CommentDataset(torch.utils.data.Dataset):
    def __init__(self, encodings, labels):
        self.encodings = encodings
        self.labels = labels

    def __getitem__(self, idx):
        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
        item['labels'] = torch.tensor(self.labels[idx])
        return item

    def __len__(self):
        return len(self.labels)


train_dataset = CommentDataset(train_token, train_label)
val_dataset = CommentDataset(val_token, val_label)
test_dataset = CommentDataset(test_token, test_label)
